import cv2
import numpy as np


def doPseudoColorEnhance(src: np.ndarray):
    """
    执行伪彩色图像增强
    :param src: 输入张量格式H*W*1或H*W，0-255 uint8 或 0-1 float
    :return: 输出张量格式H*W*3 顺序RGB 0-255 uint8
    """
    if len(src.shape) == 3:
        src = src[:, :, 0]
    if src.dtype == 'uint8':
        src = src.copy() / 255
    r = np.abs(np.sin(np.pi * src))
    g = np.abs(np.sin(np.pi * src + np.pi / 3))
    b = np.abs(np.sin(np.pi * src + 2 * np.pi / 3))
    rgb = (np.stack((r, g, b), axis=2) * 255).astype(np.uint8)
    return rgb


if __name__ == '__main__':
    '''
    伪彩色图像增强功能测试用
    '''
    # 读入图像，注意BGR
    src = cv2.imread('4PseudoColorEnhance.jpg')  # H*W*C
    src = src[:, :, [0]]
    # src = src[:, :, 0]
    print(src.shape)
    rgb = doPseudoColorEnhance(src)
    cv2.imshow("0", rgb[:, :, ::-1])
    cv2.waitKey(0)
